from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from typing import Any, List, Optional
import requests
from llama_index.core.bridge.pydantic import Field, PrivateAttr
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core import Settings
from llama_index.llms.deepseek import DeepSeek
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.storage.docstore import SimpleDocumentStore


class SiliconFlowEmbedding(BaseEmbedding):
    model_name: str = Field(
        default="BAAI/bge-large-zh-v1.5",
        description="The name of the model to use."
    )
    api_key: str = Field(
        default="Bearer sk-ngekgvpblpqccnltgqydfzdxsktivlwmaoyjdfklsdjsadtz",
        description="The API key for the embedding service."
    )
    api_url: str = Field(
        default="https://api.siliconflow.cn/v1/embeddings",
        description="The URL for the embedding API."
    )

    def __init__(
            self,
            model_name: str = "BAAI/bge-large-zh-v1.5",
            api_key: str = "Bearer sk-ngekgvpblpqccnltgqydfzdxsktivlwmaoyjdfklsdjsadtz",
            api_url: str = "https://api.siliconflow.cn/v1/embeddings",
            **kwargs
    ):
        super().__init__(
            model_name=model_name,
            api_key=api_key,
            api_url=api_url,
            **kwargs
        )

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_embedding(text)

    def _get_query_embedding(self, query: str) -> List[float]:
        return self._get_embedding(query)

    def _get_text_embedding(self, text: str) -> List[float]:
        return self._get_embedding(text)

    def _get_embedding(self, text: str) -> List[float]:
        headers = {
            "Authorization": self.api_key,
            "Content-Type": "application/json"
        }

        payload = {
            "model": self.model_name,
            "input": text,
            "encoding_format": "float"
        }

        # print(text)
        response = requests.post(self.api_url, json=payload, headers=headers)
        response.raise_for_status()

        data = response.json()
        if "data" not in data or len(data["data"]) == 0:
            raise ValueError("Invalid response from embedding API")

        return data["data"][0]["embedding"]

    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
        return [self._get_embedding(text) for text in texts]
